"""
title: Open WebUI 0.6 Pipeline - AI Agent Database
author: Gou Feng
date: 2025-04-10
version: 1.0
license: MIT
description: 将用户消息转发给生成 SQL 语句的模型处理
requirements: requests
"""

import json
import requests
from typing import List, Optional
from pydantic import BaseModel
# from schemas import OpenAIChatMessage

class Open_WebUI_0_6_Pipeline_AI_Database_Manager:

    # 配置类 - 定义过滤器的连接目标和优先级
    class Valves(BaseModel):

        # 作用目标管道 (默认作用于 llama3:latest 模型)
        # - 格式说明: ["*"] 表示作用于所有模型, ["model1", "model2"] 指定特定模型
        pipelines: List[str] = ["*"]

        # 执行优先级（数值越小越先执行）
        priority: int = 0

    def __init__(self):

        # 初始化过滤器属性
        self.type = "filter"           # 声明为过滤器类型
        self.name = "AI 代理 - 数据库" # 中文展示名称
        self.valves = self.Valves()    # 加载配置参数

    # 生命周期钩子函数
    async def on_startup(self):
        """服务器启动时触发 (可初始化资源)"""
        print(f"[{self.name}] 过滤器启动")

    async def on_shutdown(self):
        """服务器关闭时触发 (可清理资源)"""
        print(f"[{self.name}] 过滤器关闭")

    # 核心处理逻辑
    async def inlet(self, body: dict, user: Optional[dict] = None) -> dict:
        """
        处理请求数据的主入口
        :param body: 包含消息列表的请求体 (格式: {"messages": [...]})
        :param user: 用户信息 (可选)
        :return: 修改后的请求体
        """

        # 遍历所有消息 (原始数据格式为字典列表)
        for message in body.get("messages", []):

            # 仅处理用户消息 (role == "user")
            if message.get("role") == "user":

                # 提取用户消息内容
                original_content = message.get("content", "")

                if ("数据查询" in original_content):

                    # 通过 API 调用其它模型 (已经经过 System Prompt 调校的 LLM 模型) 完成数据库操作
                    Headers   = {"Content-Type": "application/json"}
                    JSON_Data = {"args": [original_content]}
                    # ..............................
                    response = requests.post("http://localhost:5000/api/ai-agent-database", headers=Headers, json=JSON_Data)
                    # ..............................
                    JSON_Response = json.loads(response.text)
                    # ..............................
                    Model_Content = JSON_Response["content"]

                    # 将其它模型 (已经经过 System Prompt 调校的 LLM 模型) 的返回结果交由 Open WebUI 0.6 当前对话窗口的 LLM 模型返回给用户
                    modified_content = f"任务: 将用户提供的信息以最佳方式展示\n提示: 用户提供的信息将放在 [CONTEXT] 中\n开始任务:\n[CONTEXT]\n{Model_Content}"
                    message["content"] = modified_content

                    # 调试日志 (实际使用可注释掉)
                    print(f"[{self.name}] 已修改用户消息: {modified_content[:50]}...")

        # 返回修改后的完整请求体
        return body

# EOF Signed by GF.
